#include <iostream>
#include "PNP.h"

using namespace std;
using namespace cv;

VideoCapture capture;
int g_slider_position = 0;
int g_run = -1;
int g_dontset = 0;
void onTrackbarSlide(int pos, void *) //当用户拖动滑块的时候，利用函数callback（回调）出当前滑块的位置（视频帧数）
{
    capture.set(CAP_PROP_POS_FRAMES, pos);
    if (!g_dontset)
    {
        g_run = 1;
    }
    g_dontset = 0;
}

PNP::PNP()
{
    _is_recorded = false;
    _color = RED;
}

PNP::PNP(bool is_recorded = false)
{
    _is_recorded = is_recorded;
    _color = RED;
}

PNP::PNP(bool is_recorded = false, int color = RED)
{
    _is_recorded = is_recorded;
    _color = color;
}

void PNP::set_objpoint(double length, double width)
{
    _half_length = length / 2.0;
    _half_width = width / 2.0;
    _obj = vector<Point3f>{
        Point3f(_half_length, _half_width, 0),
        Point3f(_half_length, -_half_width, 0),
        Point3f(-_half_length, _half_width, 0),
        Point3f(-_half_length, -_half_width, 0)};
}

void PNP::set_camera_params(cv::Mat camera_matrix, cv::Mat distcoeffs)
{
    _camera_matrix = camera_matrix;
    _distcoeffs = distcoeffs;
}

int PNP::ImageprocessAndPnP(std::string video_path)
{
    namedWindow("效果视频", WINDOW_GUI_EXPANDED);
    capture.open(video_path);
    if (!capture.isOpened())
    {
        cout << "video cannot be opened! Please check your video path." << endl;
        return -1;
    }

    int frames = (int)capture.get(CAP_PROP_FRAME_COUNT);
    int tmpw = (int)capture.get(CAP_PROP_FRAME_WIDTH);
    int tmph = (int)capture.get(CAP_PROP_FRAME_HEIGHT);
    createTrackbar("Position", "效果视频", &g_slider_position, frames, onTrackbarSlide);

    if (_is_recorded)
    {
        double fps = capture.get(CAP_PROP_FPS);
        Size size(
            (int)capture.get(CAP_PROP_FRAME_WIDTH),
            (int)capture.get(CAP_PROP_FRAME_HEIGHT));
        writer.open("./代码效果.avi", VideoWriter::fourcc('M', 'P', 'G', '2'), fps, size);
    }

    Mat src;
    Mat gray_img;
    Mat binBright_img;
    vector<Mat> channels;
    vector<vector<Point>> lightContours; //存储若干边框像素点坐标的数组（二维数组，一个维度表示某个边框的像素点坐标，另一个维度表示某个边框）
    RotatedRect rrect;                   //边框外接矩形对象
    vector<Point2f> light_center;        //存储灯带中心坐标的数组
    vector<Point2f> light_point;         //存储灯带角点的数组
    while (capture.read(src))
    {
        if (g_run != 0)
        {
            int current_pos = (int)capture.get(CAP_PROP_POS_FRAMES); //得到目前视频帧数位置
			g_dontset = 1; //将g_dontset设置为1，使得下一个callback函数不会将系统置于单步模式
            split(src, channels);
            if (_color) //判断是否灯带颜色为红光
            {

                gray_img = channels.at(2) - channels.at(0); //增强r通道的灰度值，即增强红色光，并写入到单通道图像中
            }
            else
            {
                gray_img = channels.at(0) - channels.at(2); //增强b通道的灰度值，即增强蓝色光，并写入到单通道图像中
            }

            threshold(gray_img, binBright_img, 100, 255, THRESH_BINARY); //二值化视频帧
            // imshow("二值化视频", binBright_img);

            findContours(binBright_img.clone(), lightContours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE); //找出所有边框的像素点坐标

            for (int n = 0, cnt = 0; n < lightContours.size(); n++) //遍历所有轮廓
            {
                rrect = minAreaRect(lightContours[n]); //返回某个边框外接矩形
                Point2f cpt = rrect.center;            //外接矩形的中心点赋值给cpt

                /* 剔除不符合装甲板灯条灯光的轮廓 */
                double bigger_line;
                if (rrect.size.width < rrect.size.height)
                {
                    bigger_line = (double)rrect.size.height; //外接矩形的height较大，赋值给bigger_line
                }
                else
                {
                    bigger_line = (double)rrect.size.width; //外接矩形的width较大，赋值给bigger_line
                }
                if (bigger_line < 20) //筛掉长度过小的轮廓
                {
                    continue;
                }

                auto max_y = (double)lightContours[n][0].y;
                auto min_y = (double)lightContours[n][0].y;
                double temp;
                Point up_point = lightContours[n][0], low_point = lightContours[n][0];
                for (size_t i = 0; i < lightContours[n].size(); i++)
                {
                    temp = (double)lightContours[n][i].y;
                    if (temp > max_y)
                    {
                        max_y = temp;
                        up_point = lightContours[n][i];
                    }
                    if (temp < min_y)
                    {
                        min_y = temp;
                        low_point = lightContours[n][i];
                    }
                }
                light_point.push_back(up_point);  //存入下角点(y坐标较大的角点)
                light_point.push_back(low_point); //存入上角点(y坐标较大的角点)
                light_center.push_back(cpt);      //存入灯条外接矩形中心点
            }

            /* 标定装甲板区域 */
            if (light_point.size() > 3)
            {
                line(src, light_point[0], light_point[3], Scalar(0, 128, 0), 1, 8, 0);
                line(src, light_point[1], light_point[2], Scalar(0, 128, 0), 1, 8, 0);
                // circle(src, light_point[0], 5, Scalar(0,0,255), 1, 8 ,0);
            }

            /* 进行姿态解算，并标定装甲板中心点 */
            if (light_center.size() > 1)
            {
                auto result = p4psolution(light_point);                      //得到保存有姿态角和欧式距离的数组
                auto armor_center = (light_center[0] + light_center[1]) / 2; //装甲板中心
                circle(src, armor_center, 5, Scalar(0, 0, 255), -1, 8, 0);   //标定装甲板中心点
                string euclidean_distance = to_string(result[3]);
                euclidean_distance = euclidean_distance + "mm";
                putText(src, euclidean_distance, armor_center, 0, 0.3, Scalar(255, 255, 255), 1, 8, 0);

                string roll_err = to_string(result[0]);
                string yaw_err = to_string(result[1]);
                string pitch_err = to_string(result[2]);
                string err = "roll:" + roll_err + "  " + "yaw:" + yaw_err + "  " + "pitch:" + pitch_err;
                putText(src, err, Point(src.cols / 2 - 180, src.rows - 20), 0, 0.5, Scalar(0, 255, 0), 1, 8, 0);
                cout << err << "  "
                     << "distance:" << euclidean_distance << endl; //终端输出姿态角和距离
            }
            light_center.clear();
            light_point.clear();
            
			setTrackbarPos("Position","效果视频",current_pos);
            imshow("效果视频", src);
            if (_is_recorded)
            {
                writer << src;
            }
            g_run -= 1;
        }

        char c = (char)waitKey(10);
        if (c == ' ') //按下空格停止
        {
            if (waitKey(0) == 27)
                break; //按下esc退出
        }
        if (c == 's') // single step
		{
			g_run = 1; cout << "Single step, run = " << g_run << endl;
		}
		if (c == 'r') // run mode
		{
			g_run = -1; cout << "Run mode, run = " << g_run << endl;
		}
        
    } //while循环

    return 0;
}

vector<double> PNP::p4psolution(const vector<Point2f> objectPoints)
{
    vector<double> result;
    vector<double> result_eigen;
    Mat _rVec = Mat::zeros(3, 1, CV_64FC1); //旋转向量
    Mat _tVec = Mat::zeros(3, 1, CV_64FC1); //平移向量
    solvePnP(_obj, objectPoints, _camera_matrix, _distcoeffs, _rVec, _tVec);
    double euclidean_distance = sqrt(_tVec.at<double>(0, 0) * _tVec.at<double>(0, 0) +
                                     _tVec.at<double>(1, 0) * _tVec.at<double>(1, 0) + _tVec.at<double>(2, 0) * _tVec.at<double>(2, 0)); //计算三维空间下的欧式距离
    Rodrigues(_rVec, _rVec);         //罗德里格斯变换得到等价旋转矩阵                                                                                                    //由旋转向量得到等价旋转矩阵

    /* Eigen库求解欧拉角 */
    Eigen::Matrix3f rVec;
    cv2eigen(_rVec, rVec);
    Eigen::Vector3f eular_angle = rVec.eulerAngles(2, 1, 0); //由旋转矩阵rVec得到欧拉角(zyx)

    /* 旋转矩阵元素求解欧拉角 */
    double rollErr = atan2(_rVec.at<double>(1, 0), _rVec.at<double>(0, 0)) / CV_PI * 180; //z轴旋转角roll
    double yawErr = atan2(-_rVec.at<double>(2, 0), sqrt(pow(_rVec.at<double>(2, 0), 2) +
                                                        pow(_rVec.at<double>(2, 2), 2))) /
                    CV_PI * 180;                                                           //y轴旋转角yaw
    double pitchErr = atan2(_rVec.at<double>(2, 1), _rVec.at<double>(2, 2)) / CV_PI * 180; //x轴旋转角pitch

    result.resize(4);
    result[0] = rollErr;            //z轴旋转角roll
    result[1] = yawErr;             //y轴旋转角yaw
    result[2] = pitchErr;           //x轴旋转角pitch
    result[3] = euclidean_distance; //欧式距离

    result_eigen.resize(4);
    result_eigen[0] = eular_angle[0] / M_PI * 180; //图像z轴旋转角
    result_eigen[1] = eular_angle[1] / M_PI * 180; //图像y轴旋转角
    result_eigen[2] = eular_angle[2] / M_PI * 180; //图像x轴旋转角
    result_eigen[3] = euclidean_distance;

    return result_eigen;
}

PNP::~PNP()
{
    capture.release();
    writer.release();
}